import torch
import numpy as np 
import torch.nn as nn 
import math
import torch.nn.functional as F 
from PIL import Image 
from torch.autograd import Variable
from torch.utils.data import DataLoader 
from torch.utils.data.dataset import Dataset 
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
from random import shuffle 


class MyDataset(Dataset):
    def __init__(self, train_lines, image_size):
        self.train_lines = train_lines
        self.image_size = image_size
        self.train_batches = len(train_lines)
    
    def __len__(self):
        return self.train_batches
    
    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def get_data(self, lines, input_shape):
        jitter=.3
        hue=.1
        sat=1.5
        val=1.5
        line = lines.split()
        image = Image.open(line[0])
        iw, ih = image.size
        h, w = input_shape
        box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])

        # 调整图片大小
        new_ar = w / h * self.rand(1 - jitter, 1 + jitter) / self.rand(1 - jitter, 1 + jitter)
        scale = self.rand(.25, 2)
        if new_ar < 1:
            nh = int(scale * h)
            nw = int(nh * new_ar)
        else:
            nw = int(scale * w)
            nh = int(nw / new_ar)
        image = image.resize((nw, nh), Image.BICUBIC)

        # 放置图片
        dx = int(self.rand(0, w - nw))
        dy = int(self.rand(0, h - nh))
        new_image = Image.new('RGB', (w, h),
                              (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255)))
        new_image.paste(image, (dx, dy))
        image = new_image

        # 是否翻转图片
        flip = self.rand() < .5
        if flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)

        # 色域变换
        hue = self.rand(-hue, hue)
        sat = self.rand(1, sat) if self.rand() < .5 else 1 / self.rand(1, sat)
        val = self.rand(1, val) if self.rand() < .5 else 1 / self.rand(1, val)
        x = rgb_to_hsv(np.array(image) / 255.)
        x[..., 0] += hue
        x[..., 0][x[..., 0] > 1] -= 1
        x[..., 0][x[..., 0] < 0] += 1
        x[..., 1] *= sat
        x[..., 2] *= val
        x[x > 1] = 1
        x[x < 0] = 0
        image_data = hsv_to_rgb(x) * 255  # numpy array, 0 to 1

        # 调整目标框坐标
        box_data = np.zeros((len(box), 5))
        if len(box) > 0:
            np.random.shuffle(box)
            box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
            box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
            if flip:
                box[:, [0, 2]] = w - box[:, [2, 0]]
        if len(box) == 0:
            return image_data, []
        if (box[:, :4] > 0).any():
            return image_data, box
        else:
            return image_data, []


    def __getitem__(self, index):
        lines = self.train_lines
        img, box = self.get_data(lines[index], self.image_size)
        boxes = np.array(box[:,:4],dtype=np.float32)
        boxes[:,0] = boxes[:,0]/self.image_size[1]
        boxes[:,1] = boxes[:,1]/self.image_size[0]
        boxes[:,2] = boxes[:,2]/self.image_size[1]
        boxes[:,3] = boxes[:,3]/self.image_size[0]
        box = np.concatenate([boxes,box[:,-1:]],axis=-1)
        img = np.array(img, dtype=np.float32)
        img = img / 255
        img  = np.transpose(img,(2,0,1))
        return img, box